from typing import Annotated
from langchain_core.tools import tool, InjectedToolCallId
from langgraph.prebuilt import create_react_agent, InjectedState
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.types import Command
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage, AIMessage

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

def transfer_to_hotel_assistant(state: Annotated[MessagesState, InjectedState]):
    """Transfer to hotel assistant."""
    return Command(
        # name of the agent (node) to go to
        goto="hotel_assistant",
        # data to send to the agent
        # update={"messages": state["messages"] + [ToolMessage(content="Successfully transferred to hotel assistant", tool_call_id="12345")]},
        update={"messages": [AIMessage(content="Successfully transferred to hotel assistant")]}
        # indicate to LangGraph that we need to navigate to
        # agent node in a parent graph
        # graph=Command.PARENT,
    )

def transfer_to_flight_assistant(state: MessagesState):
    """Transfer to flight assistant."""
    return Command(
        # name of the agent (node) to go to
        goto="flight_assistant",
        # data to send to the agent
        # update={"messages": [state["messages"]]},
        # indicate to LangGraph that we need to navigate to
        # agent node in a parent graph
        # graph=Command.PARENT,
    )

# Simple agent tools
def book_hotel(hotel_name: str):
    """Book a hotel"""
    return f"Successfully booked a stay at {hotel_name}."

def book_flight(from_airport: str, to_airport: str):
    """Book a flight"""
    return f"Successfully booked a flight from {from_airport} to {to_airport}."

# Define agents
flight_assistant = create_react_agent(
    model=llm,
    # tools=[book_flight, transfer_to_hotel_assistant],
    tools=[book_flight],
    prompt="You are a flight booking assistant",
    name="flight_assistant"
)
hotel_assistant = create_react_agent(
    model=llm,
    # tools=[book_hotel, transfer_to_flight_assistant],
    tools=[book_hotel],
    prompt="You are a hotel booking assistant",
    name="hotel_assistant"
)

# Define multi-agent graph
multi_agent_graph = (
    StateGraph(MessagesState)
    .add_node(flight_assistant)
    .add_node(hotel_assistant)
    .add_node(transfer_to_hotel_assistant)
    .add_edge(START, "flight_assistant")
    .add_edge("flight_assistant", 'transfer_to_hotel_assistant')
    .compile()
)

# Run the multi-agent graph
for chunk in multi_agent_graph.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "Please book a flight from BOS to JFK and a stay at McKittrick Hotel"
            }
        ]
    }
):
    print("------------------------------------------------------------------------------------")
    if "flight_assistant" in chunk and chunk["flight_assistant"] and chunk["flight_assistant"]["messages"]:
        print("--------------------------------  flight_assistant   -------------------------------")
        messages = chunk["flight_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    elif "transfer_to_hotel_assistant" in chunk and chunk["transfer_to_hotel_assistant"] and chunk["transfer_to_hotel_assistant"]["messages"]:
        print("--------------------------------  transfer_to_hotel_assistant   -------------------------------")
        messages = chunk["transfer_to_hotel_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    elif "hotel_assistant" in chunk and chunk["hotel_assistant"] and chunk["hotel_assistant"]["messages"]:
        print("--------------------------------  hotel_assistant   -------------------------------")
        messages = chunk["hotel_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    else:
        print(chunk)